"""
Copyright (c) 2022 Huawei Technologies Co.,Ltd.

openGauss is licensed under Mulan PSL v2.
You can use this software according to the terms and conditions of the Mulan PSL v2.
You may obtain a copy of Mulan PSL v2 at:

          http://license.coscl.org.cn/MulanPSL2

THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
See the Mulan PSL v2 for more details.
"""
"""
Case Type   : python驱动psycopg2
Case Name   : python驱动psycopg2公共函数
Create At   : 2023-08-23
Owner       : opentestcase022
Description :
    1、python驱动psycopg2公共函数
Expect      :
    1、python驱动psycopg2公共函数
History     :
"""
import os
import subprocess
import traceback

from yat.test import macro

from testcase.utils.CommonSH import CommonSH
from testcase.utils.Constant import Constant
from testcase.utils.Logger import Logger


class Psycopg2Common:
    def __init__(self, **kwargs):
        self.log = Logger()
        self.pri_sh = CommonSH('PrimaryDbUser')
        self.cons = Constant()
        self.pkg_path = os.path.dirname(macro.DB_SCRIPT_PATH)
        self.new_dir = "opg_psycopg2"
        self.cur_ip = subprocess.getoutput("hostname -I|awk '{{print $1}}'")
        self.dbname = kwargs['dbname']
        self.user = kwargs['user']
        self.password = kwargs['password']
        self.host = kwargs['host']
        self.port = kwargs['port']
        self.py_path = None

    def decompress_pkg(self):
        """
        解压***-Python.tar.gz得到lib和psycopg2
        :return: True执行成功，False执行失败
        """
        text = '----解压***-Python.tar.gz得到lib----'
        self.log.info(text)
        cmd = f"cd {self.pkg_path} && " \
            f"if [ -e opg_psycopg2 ]; then echo 'exists'; " \
            f"else mkdir opg_psycopg2; fi && " \
            f"if [ -e openGauss-*-Python.tar.gz ]; " \
            f"then tar -zxf openGauss-*-Python.tar.gz -C {self.new_dir} && " \
            f"ls -l {self.new_dir}; " \
            f"else echo 'PDBC package not exists'; " \
            f"fi "
        self.log.info(cmd)
        res = subprocess.getoutput(cmd)
        self.log.info(res)
        if 'lib' in res:
            return True
        else:
            return False

    def set_new_psycopg2(self):
        """
        使用root用户将psycopg2拷贝到python安装目录下的site-packages文件夹下，并修改psycopg2目录权限为755
        :return:True执行成功，False执行失败
        """
        flag = False
        cmd = "find /usr/lib -type d -name 'site-packages'|grep python3|" \
              "head -1"
        self.log.info(cmd)
        self.py_path = subprocess.getoutput(cmd)
        self.log.info(self.py_path)

        if self.py_path.strip():
            pdbc_path = os.path.join(self.pkg_path, self.new_dir, 'psycopg2')
            cmd = f"\\cp -r {pdbc_path} {self.py_path}/ && " \
                f"chmod 755 -R {os.path.join(self.py_path, 'psycopg2')}"
            self.log.info(cmd)
            res = subprocess.getoutput(cmd)
            self.log.info(res)
            flag = True
        return flag

    def set_lib_path(self):
        """
        将解压后的lib目录，配置在LD_LIBRARY_PATH中
        :return:True执行成功，False执行失败
        """
        text = '----配置lib,连接数据库----'
        self.log.info(text)
        lib_path = os.path.join(self.pkg_path, self.new_dir, "lib")
        os.environ['LD_LIBRARY_PATH'] += f':{lib_path}'
        os.environ['PYTHONPATH'] += f':{self.py_path}'
        cmd = f'env|grep LD_LIBRARY_PATH|grep "{lib_path}";' \
            f'echo $PYTHONPATH|grep "{self.py_path}"'
        self.log.info(cmd)
        res = subprocess.getoutput(cmd)
        self.log.info(res)
        if not res.strip():
            return False
        else:
            return True

    def bak_conf(self, bak_flag=True, f_name=None):
        """
        备份/还原pg_hba文件
        :param bak_flag: True备份，False还原
        :param f_name: 默认为pg_hba
        :return:True执行成功，False执行失败
        """
        is_right = False
        if not f_name:
            f_name = os.path.join(macro.DB_INSTANCE_PATH,
                                  macro.PG_HBA_FILE_NAME)
        else:
            f_name = os.path.join(macro.DB_INSTANCE_PATH, f_name)
        if bak_flag:
            cmd = f"\\cp {f_name} {f_name}_bak.bak"
        else:
            cmd = f"if [ -e {f_name}_bak.bak ]; " \
                f"then \\cp {f_name}_bak.bak {f_name}; " \
                f"rm -rf {f_name}_bak.bak; fi"
        self.log.info(cmd)
        res = self.pri_sh.node.sh(cmd).result()
        self.log.info(res)
        if not bak_flag and not res.strip():
            res1 = self.pri_sh.reload_db_config()
            if 'server signaled' in res1:
                is_right = True
        elif bak_flag and not res.strip():
            is_right = True
        return is_right

    def set_pghba(self, open_flag=True, db=None, user=None, mode='sha256'):
        """
        开启/取消白名单
        :param open_flag: True开启，False取消
        :param db: 数据库名
        :param user: 数据库用户名
        :param mode: 认证方式
        :return: True执行成功，False执行失败
        """
        pghba_path = os.path.join(macro.DB_INSTANCE_PATH,
                                  macro.PG_HBA_FILE_NAME)
        hostmode = 'host' if mode != 'cert' else 'hostssl'
        grep_cmd = f"grep {self.cur_ip} {pghba_path}|grep sha256|" \
            f"grep -v '#'|awk '{{{{print $2, $3}}}}'"
        param_list = list()
        set_mode = mode if open_flag else ''
        db_name = db if db else 'all'
        db_user = user if user else 'all'
        try:
            if mode == 'cert':
                pghba_param = f'{hostmode} {db_name} {db_user} ' \
                    f'{self.cur_ip}/32 {set_mode}'
                param_list.append(pghba_param)
                pghba_param = f'{hostmode} all all 127.0.0.1/32 {set_mode}'
                param_list.append(pghba_param)
            elif db or user:
                pghba_param = f'{hostmode} {db_name} {db_user} ' \
                    f'{self.cur_ip}/32 {set_mode}'
                param_list.append(pghba_param)
            else:
                grep_res = self.pri_sh.node.sh(grep_cmd).result()
                self.log.info(grep_res)
                if grep_res.strip():
                    for line in grep_res.strip().splitlines():
                        db_name = db if db else line.strip().split()[0]
                        db_user = user if user else line.strip().split()[1]
                        pghba_param = f'{hostmode} {db_name} {db_user} ' \
                            f'{self.cur_ip}/32 {set_mode}'
                        param_list.append(pghba_param)
            if param_list:
                self.log.info(param_list)
                str_params = ' '.join([f'-h "{i}"' for i in param_list])
                guc_cmd = f'source {macro.DB_ENV_PATH};' \
                    f'gs_guc reload -D {macro.DB_INSTANCE_PATH} ' \
                    f'{str_params}'
                self.log.info(guc_cmd)
                guc_res = self.pri_sh.node.sh(guc_cmd).result()
                self.log.info(guc_res)
                assert 'Failed instances: 0' in guc_res
            else:
                raise Exception('NotFound: 未找到记录，请检查pg_hba文件配置！')
        except Exception as e:
            grep_cmd = f"grep {self.cur_ip} {pghba_path}"
            grep_res = self.pri_sh.node.sh(grep_cmd).result()
            self.log.info(grep_res)
            self.log.error(f'str(e): {e}')
            self.log.error(f'repr(e): {repr(e)}')
            self.log.error(f'traceback.format_exc():\n'
                           f'{traceback.format_exc()}')
            return False
        else:
            grep_cmd = f"grep {self.cur_ip} {pghba_path}"
            grep_res = self.pri_sh.node.sh(grep_cmd).result()
            self.log.info(grep_res)
            return True

    @staticmethod
    def format_sql(sql, *args):
        """
        格式化sql语句，返回psycopg2可执行SQL语句
        :param sql: SQL语句
        :param args: SQL中传递参数序列
        :return: 返回psycopg2可执行SQL语句
        """
        if args:
            f_sql = f"cur.execute(\\\"{sql}\\\",{args})\n"
        else:
            f_sql = ''
            for s in sql.strip(';').split(';'):
                f_sql += f"cur.execute(\\\"{s}\\\")\n"
        return f_sql

    @staticmethod
    def format_sql_with_executemany(sql, *args):
        """
        executemany格式化sql语句，返回psycopg2可执行SQL语句
        :param sql: SQL语句
        :param args: SQL中传递参数序列
        :return: 返回psycopg2可执行SQL语句
        """
        f_sql = f"cur.executemany(\\\"{sql}\\\",{args})\n"
        return f_sql

    @staticmethod
    def join_run_sql(sql, rollback=None):
        """
        组合可执行的sql语句
        :param sql: sql语句
        :param rollback: None不回滚，implicit隐式回滚，explicit显式回滚
        :return:可执行的sql语句
        """
        f_sql = "cur=conn.cursor()\n"
        f_sql += f"{sql}\n"
        if not rollback:
            f_sql += "conn.commit()\n"
        elif rollback == 'explicit':
            f_sql += "conn.rollback()\n"
        if 'select' in sql.lower():
            f_sql += "results=cur.fetchall()\n"
            f_sql += "print(results)\n"
        f_sql += "cur.close()\n"
        return f_sql

    def run_sql(self, sql, by='default'):
        """
        执行sql语句，返回结果
        :param sql: sql语句
        :param by: 数据库连接方式
        :return: 连接实例
        """
        res = self.conn_db(by, sql)
        return res

    def conn_db(self, by='default', sql='', cert_path=None, ssl_mode=None):
        """
        连接数据库
        :param by:数据库连接方式
        :param sql:sql语句
        :param cert_path:by为ssl时，指定证书路径
        :param ssl_mode:ssl连接时sslmode值
        :return: 连接实例
        """
        if by == 'long_string':
            conn_info = f"\\\"dbname={self.dbname} " \
                f"user={self.user} " \
                f"password={self.password} " \
                f"host={self.host} " \
                f"port={self.port}\\\""
        elif by == 'ssl':
            conn_info = f"dbname=\\\"{self.dbname}\\\"," \
                f"user=\\\"{self.user}\\\"," \
                f"password=\\\"{self.password}\\\"," \
                f"host=\\\"{self.host}\\\"," \
                f"port=\\\"{self.port}\\\"," \
                f"sslmode=\\\"{ssl_mode}\\\"," \
                f"sslcert=\\\"{os.path.join(cert_path, 'client.crt')}\\\"," \
                f"sslkey=\\\"{os.path.join(cert_path, 'client.key')}\\\", " \
                f"sslrootcert=\\\"{os.path.join(cert_path, 'cacert.pem')}\\\""
        else:
            conn_info = f"dbname=\\\"{self.dbname}\\\"," \
                f"user=\\\"{self.user}\\\"," \
                f"password=\\\"{self.password}\\\"," \
                f"host=\\\"{self.host}\\\"," \
                f"port=\\\"{self.port}\\\""
        cmd = f"python3 -c \"import psycopg2\n" \
            f"conn = psycopg2.connect({conn_info})\n" \
            f"print(conn.dsn)\n" \
            f"{sql}\n" \
            f"conn.close()\""
        scp_cmd = f'''expect <<EOF
            set timeout 300
            spawn {cmd}
            expect {{
                "*Enter PEM pass phrase:*" \
                {{ send "{self.pri_sh.node.ssh_password}\\n" }}
                eof {{ send_user "exit\\n" }}
            }}\n''' + 'EOF'
        self.log.info(scp_cmd)
        res = subprocess.getoutput(scp_cmd)
        self.log.info(res)
        return res

    def clean_decompress_file(self):
        """
        清理解压缩文件
        :return:True执行成功，False执行失败
        """
        cmd = f"if [ -e {os.path.join(self.pkg_path, self.new_dir)} ]; " \
            f"then rm -rf {os.path.join(self.pkg_path, self.new_dir)}; " \
            f"fi "
        self.log.info(cmd)
        res = subprocess.getoutput(cmd)
        self.log.info(res)
        if not res.strip():
            return True
        else:
            return False
